import json
import os
import cv2

# root path, including images filefolder, annos.txt(bbox annotation), classes.txt(class label),
# and annotations filefolder(if not existed, then it will be created automatically, to save the final json file
root_path = '/home/v/project/TestCoco'
# to create train data or validation data
phase = 'val'
# split thresh for train and val data
split = 8000

# open class label
with open(os.path.join(root_path, 'classes.txt')) as f:
   classes = f.read().strip().split()

'''
indexes = [f for f in os.listdir(os.path.join(root_path, 'images'))]
for k, index in enumerate(indexes):
    im = cv2.imread(os.path.join(root_path, 'images/') + index)
    height, width, _ = im.shape
    print(height, width)
    cv2.imshow(index, im)
    cv2.waitKey(-1)
cv2.destroyAllWindows()
'''

'''
# create the correspondences between class label and digital id
for i, cls in enumerate(classes, 1):
    dataset['categories'].append({'id': i, 'name': cls, 'supercategory': 'mark'})

# read picture names in the images filefolder
indexes = [f for f in os.listdir(os.path.join(root_path, 'images'))]

# judge whether to create train data or val data
if phase == 'train':
    indexes = [line for i, line in enumerate(_indexes) if i <= split]
elif phase == 'val':
    indexes = [line for i, line in enumerate(_indexes) if i > split]

# read bbox info
with open(os.path.join(root_path, 'annos.txt')) as tr:
    annos = tr.readlines()

for k, index in enumerate(indexes):
    # read image using opencv, to obtain the width and height of the images
    im = cv2.imread(os.path.join(root_path, 'images/') + index)
    height, width, _ = im.shape

    # add image info to dataset
    dataset['images'].append({'file_name': index,
                              'id': k,
                              'width': width,
                              'height': height})

for ii, anno in enumerate(annos):
    parts = anno.strip().split()

    # if image name corresponds to the annotation name, then add the annotation
    if parts[0] == index:
        # classes
        cls_id = parts[1]
        # x_min
        x1 = float(parts[2])
        # y_min
        y1 = float(parts[3])
        # x_max
        x2 = float(parts[4])
        # y_max
        y2 = float(parts[5])
        width = max(0, x2 - x1)
        height = max(0, y2 - y1)
        dataset['annotations'].append({
            'area': width * height,
            'bbox': [x1, y1, width, height],
            'category_id': int(cls_id),
            'id': i,
            'image_id': k,
            'iscrowd': 0,
            # mask, rectangle that counts four points from the left up corner point counterwisely
            'segmentation': [[x1, y1, x2, y1, x2, y2, x1, y2]]
        })

# filefolder to save the result
folder = os.path.join(root_path, 'annotations')
if not os.path.exists(folder):
    os.makedirs(folder)
json_name = os.path.join(root_path, 'annotations/{}.json'.format(phase))
with open(json_name, 'w') as f:
    json.dump(dataset, f)
'''
